{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FFI 容器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm.ffi as tvm_ffi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数组"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = tvm_ffi.convert([1, 2, 3])\n",
    "assert isinstance(a, tvm_ffi.Array)\n",
    "assert len(a) == 3\n",
    "assert a[-1] == 3\n",
    "a_slice = a[-3:-1]\n",
    "assert (a_slice[0], a_slice[1]) == (1, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `map`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm import te\n",
    "a = te.var(\"a\")\n",
    "b = te.var(\"b\")\n",
    "amap = tvm_ffi.convert({a: 2, b: 3})\n",
    "assert a in amap\n",
    "assert len(amap) == 2\n",
    "dd = dict(amap.items())\n",
    "assert a in dd\n",
    "assert b in dd\n",
    "assert a + 1 not in amap\n",
    "assert {x for x in amap} == {a, b}\n",
    "assert set(amap.keys()) == {a, b}\n",
    "assert set(amap.values()) == {2, 3}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数组映射"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = tvm_ffi.convert([[1, 2, 3], {\"A\": 5, \"B\": 6}])\n",
    "assert isinstance(a, tvm_ffi.Array)\n",
    "assert len(a) == 2\n",
    "assert isinstance(a[0], tvm_ffi.Array)\n",
    "assert isinstance(a[1], tvm_ffi.Map)\n",
    "assert tuple(a[0]) == (1, 2, 3)\n",
    "assert a[1][\"A\"] == 5\n",
    "assert a[1][\"B\"] == 6\n",
    "\n",
    "a = tvm_ffi.convert([1, 2, 3])\n",
    "assert dir(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = te.var(\"a\")\n",
    "b = te.var(\"b\")\n",
    "amap = tvm_ffi.convert({a: 2, b: 3})\n",
    "assert dir(amap)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = te.var(\"a\")\n",
    "b = te.var(\"b\")\n",
    "amap = tvm_ffi.convert({a: 2, b: 3})\n",
    "assert isinstance(amap, tvm.ffi.Map)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 整数映射"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "amap = tvm_ffi.convert({3: 2, 4: 3})\n",
    "assert 3 in amap\n",
    "assert len(amap) == 2\n",
    "dd = dict(amap.items())\n",
    "assert 3 in dd\n",
    "assert 4 in dd\n",
    "assert 5 not in amap\n",
    "assert tuple(amap.items()) == ((3, 2), (4, 3))\n",
    "assert tuple(amap.keys()) == (3, 4)\n",
    "assert tuple(amap.values()) == (2, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 字符串映射"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = []\n",
    "for i in reversed(range(10)):\n",
    "    data.append((f\"a{i}\", i))\n",
    "amap = tvm_ffi.convert({k: v for k, v in data})\n",
    "assert tuple(amap.items()) == tuple(data)\n",
    "for k, v in data:\n",
    "    assert k in amap\n",
    "    assert amap[k] == v\n",
    "    assert amap.get(k) == v\n",
    "\n",
    "assert tuple(k for k in amap) == tuple(k for k, _ in data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## KeyError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytest\n",
    "amap = tvm_ffi.convert({3: 2, 4: 3})\n",
    "with pytest.raises(KeyError):\n",
    "    amap[5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## {data}`repr`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = tvm_ffi.convert([1, 2, 3])\n",
    "assert str(a) == \"[1, 2, 3]\"\n",
    "amap = tvm_ffi.convert({3: 2, 4: 3})\n",
    "assert str(amap) == \"{3: 2, 4: 3}\"\n",
    "\n",
    "smap = tvm_ffi.convert({\"a\": 1, \"b\": 2})\n",
    "assert str(smap) == \"{'a': 1, 'b': 2}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 序列化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "a = tvm_ffi.convert([1, 2, 3.5, True])\n",
    "json_str = tvm.ir.save_json(a)\n",
    "a_loaded = tvm.ir.load_json(json_str)\n",
    "assert a_loaded[1] == 2\n",
    "assert a_loaded[2] == 3.5\n",
    "assert a_loaded[3] == True\n",
    "assert isinstance(a_loaded[3], bool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = te.var(\"a\")\n",
    "b = te.var(\"b\")\n",
    "amap = tvm_ffi.convert({a: 2, b: 3})\n",
    "json_str = tvm.ir.save_json(amap)\n",
    "amap = tvm.ir.load_json(json_str)\n",
    "assert len(amap) == 2\n",
    "dd = {kv[0].name: kv[1] for kv in amap.items()}\n",
    "assert dd == {\"a\": 2, \"b\": 3}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `in` 容器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = tvm_ffi.convert([\"a\", \"b\", \"c\"])\n",
    "assert \"a\" in arr\n",
    "assert tvm.tir.StringImm(\"a\") in arr\n",
    "assert \"d\" not in arr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `ndarray` 容器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tvm.nd.array([1, 2, 3])\n",
    "arr = tvm_ffi.convert([x, x])\n",
    "assert arr[0].same_as(x)\n",
    "assert arr[1].same_as(x)\n",
    "assert isinstance(arr[0], tvm.nd.NDArray)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `ReturnsVariant`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = tvm.get_global_func(\"testing.ReturnsVariant\")\n",
    "res_even = func(42)\n",
    "assert isinstance(res_even, tvm.tir.IntImm)\n",
    "assert res_even == 21\n",
    "\n",
    "res_odd = func(17)\n",
    "assert res_odd == \"argument was odd\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `AcceptsVariant`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = tvm.get_global_func(\"testing.AcceptsVariant\")\n",
    "\n",
    "assert func(\"string arg\") == \"ffi.String\"\n",
    "assert func(17) == \"ir.IntImm\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = tvm.get_global_func(\"testing.AcceptsVariant\")\n",
    "float_arg = tvm.tir.FloatImm(\"float32\", 0.5)\n",
    "\n",
    "with pytest.raises(Exception):\n",
    "    func(float_arg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
